from modelscope import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype="auto", device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

messages = [
    {
        "role": "system",
        "content": "你是研究助手，我会给你一个研究方向，你可以通过下面的方式询问我来获取相关信息。当你输出 #SEARCH(搜索字段)，时，我会把和搜索字段相关的内容告诉你。你可以自行决定合适结束，结束时输出报告，并在末尾加上 #END",
    }
]

while True:
    # 读取用户输入
    talk = input(">")
    messages.append({"role": "user", "content": talk})
    
    # llm
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(**model_inputs, max_new_tokens=512)
    generated_ids = [
        output_ids[len(input_ids) :]
        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    response: str = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(response)

    if response.strip().endswith("#END"):
        break

